from sklearn.datasets import load_iris
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score, validation_curve
import matplotlib.pyplot as plt

x, y = load_iris(return_X_y=True)

# acc = []
# for i in range(1, 11):
#     knn = KNeighborsClassifier(n_neighbors=i)
#     acc.append(cross_val_score(knn, x, y, cv=5).mean())
#
# plt.plot(range(1, 11), acc, 'o-')
# plt.show()

knn = KNeighborsClassifier()
train_acc, test_acc = validation_curve(knn, x, y, param_name='n_neighbors',
                                     param_range=range(1, 11), cv=5)
plt.plot(range(1, 11), train_acc.mean(axis=1), 'o-', label='train_acc')
plt.plot(range(1, 11), test_acc.mean(axis=1), 'o-', label='test_acc')
plt.legend()
plt.show()

